// "Copyright [2021] <Copyright Shengkai Lin>"

#ifndef CLUSTER_TASK_SCHEDULER_TASK_ALLOCATOR_H_
#define CLUSTER_TASK_SCHEDULER_TASK_ALLOCATOR_H_

#include "cluster/basic/ai_sim.h"
#include "cluster/task_scheduler/task_type/allreduce.h"
#include "cluster/network_controller/routing_algorithm/route.h"
#include "cluster/dispatcher/flowparser.h"
namespace ai_sim {

class TaskAllocator : public omnetpp::cSimpleModule {
 private:
    AllreduceAlgorithm *allreduce_algo;
    RouteAlgorithm *route_algo;
    TopoManager *topos;

 public:
    void initialize() override;
    void gen_comm_pair(taskboard&, AllreduceAlgorithm*);
    void gen_flows(taskboard&, AllreduceAlgorithm*);
    void gen_txt_flow(map<int, FlowVector>& tx_flows,
      map<int, FlowVector>& rx_flows,
      map<int, RouteVec>& routes, DetailFlow& flow);
};
}  // namespace ai_sim

#endif  // CLUSTER_TASK_SCHEDULER_TASK_ALLOCATOR_H_
